import torch
import torch.nn as nn


class LinearModel(nn.Module):
    def __init__(self, output_num=5):
        super(LinearModel, self).__init__()
        self.fc1 = nn.Linear(25, 128)
        self.fc2 = nn.Linear(128, output_num)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)
        x = self.fc2(x)
        return torch.nn.functional.softmax(x, dim=1)


if __name__ == "__main__":
    model = LinearModel()
    input = torch.randn(4, 1, 25)
    output = model(input)
    print(output)
